import numpy as np
import matplotlib.pyplot as plt
from skimage.io import imread
from skimage.filters import sobel, gaussian, median
from skimage.feature import canny
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from joblib import dump, load
from tqdm import tqdm
import pickle
import time
from skimage.color import rgb2gray  # 导入灰度转换函数

def create_clf():
    # 加载图像和分区标签
    image = imread('Sandstone_1.tif')  # 砂岩图像
    segmentation = imread('Sandstone_1_segment.tif')  # 分区标签图像

    # 确保图像是RGB格式（如果是灰度图，转换为RGB）
    if image.ndim == 2:  # 如果是灰度图
        image = np.repeat(image[:, :, np.newaxis], 3, axis=2)  # 将灰度图扩展为RGB

    # 获取图像尺寸
    print(f"图像形状: {image.shape}")

    # 提取RGB通道
    R = image[:, :, 0]  # 红色通道
    G = image[:, :, 1]  # 绿色通道
    B = image[:, :, 2]  # 蓝色通道

    # 应用不同的滤波器
    # Sobel滤波器（边缘检测）
    sobel_R = sobel(R)
    sobel_G = sobel(G)
    sobel_B = sobel(B)

    # 高斯滤波器（平滑处理）
    gaussian_R = gaussian(R, sigma=1)
    gaussian_G = gaussian(G, sigma=1)
    gaussian_B = gaussian(B, sigma=1)

    # 均值滤波器（平滑处理）
    median_R = median(R)
    median_G = median(G)
    median_B = median(B)

    # Canny滤波器（边缘检测）
    canny_R = canny(R)
    canny_G = canny(G)
    canny_B = canny(B)

    # 将所有特征合并成一个X矩阵
    # 每个像素点有12个特征：原始RGB通道 + 4个滤波器处理的通道
    X = np.stack([R, G, B, sobel_R, sobel_G, sobel_B, 
                  gaussian_R, gaussian_G, gaussian_B, 
                  median_R, median_G, median_B, 
                  canny_R, canny_G, canny_B], axis=-1)
    
    X = X.reshape(-1, 15)  # 每个像素点有15个特征

    # 获取标签y（flatten标签图像）
    # 如果标签是RGB图像，则转换为灰度图
    if segmentation.ndim == 3:  # 如果是RGB图像
        print("标签图像是RGB图像，正在转换为灰度图像...")
        segmentation = rgb2gray(segmentation)  # 转换为灰度图像
        segmentation = (segmentation * 255).astype(np.uint8)  # 转换为8位灰度图
    
    y = segmentation.flatten()

    print(f"标签图像的形状: {segmentation.shape}")

    print("完成从砂岩截面图1及其对应分区中获取X和y")
    print(f"X shape: {X.shape}")
    print(f"y shape: {y.shape}")

    # 确保X和y的形状一致
    if X.shape[0] != y.shape[0]:
        raise ValueError(f"特征和标签的样本数量不一致: X={X.shape[0]}, y={y.shape[0]}")

    # 划分训练集和测试集
    print("完成train_test_split")
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 创建并训练随机森林分类器（使用带进度条的自定义RandomForest）
     # 初始化分类器
    clf = RandomForestClassifier(n_estimators=50, max_depth=10, random_state=42, verbose=1, n_jobs=-1)


    # 使用 tqdm 包裹训练过程
    with tqdm(total=100, desc="Training Random Forest") as pbar:
        for i in range(100):
            clf.fit(X_train, y_train)
            pbar.update(1)

    print("完成随机森林模型clf的训练")

    # 输出准确率
    accuracy = clf.score(X_test, y_test)
    print(f"准确率: {accuracy:.4f}")
    

    # 保存训练好的模型
    dump(clf, 'clf.joblib')
    print("已保存随机森林模型clf到硬盘")

if __name__ == "__main__":
    create_clf()

